#pragma once

#include <memory>
#include <vector>
#include <string>
#include "blob.h"

namespace vkinfer
{
// #define get_op_type(name) OP_TYPE_##name
    enum class OpType
    {
        OP_TYPE_INVALID,
        OP_TYPE_ADD,
        OP_TYPE_MUL,
        OP_TYPE_LINEAR,
        OP_TYPE_RELU,
        OP_TYPE_CONV2D,
        OP_TYPE_SOFTMAX,
        OP_TYPE_SIGMOID,
        OP_TYPE_MAXPOOL,
        OP_TYPE_AVERAGEPOOL,
        OP_TYPE_GLOBALAVERAGEPOOL,
        OP_TYPE_FLATTEN,
        OP_TYPE_SUB,
        OP_TYPE_TRANSPOSE,
        OP_TYPE_SLICE,
        OP_TYPE_CONCAT,
    };

    OpType op_string2type(const std::string& op_type_name);
    
    struct Operator
    {
    public:
        std::string name;
        OpType type;
        std::vector<std::string> input_names;
        std::vector<std::string> output_names;
        bool top_layer = false;  // layer with no prior dependencies.
    public:
        virtual void set_weight(const std::string& name, const std::shared_ptr<Blob>& blob) {}
        virtual void forward(const std::vector<std::shared_ptr<Blob>>& inputs, 
            std::vector<std::shared_ptr<Blob>>& outputs) = 0;

        std::vector<std::string> get_input_names() const { return input_names; }
        std::vector<std::string> get_output_names() const { return output_names; }
        std::string get_name() const { return name; }
    };
}
